#!/bin/bash
#
# Copyright(C) 2021, Stamus Networks
# Written by Raphaël Brogat <rbrogat@stamus-networks.com>
#
# This script comes with ABSOLUTELY NO WARRANTY!
#
# GPLv3 <http://www.gnu.org/licenses/>.



############################################################
#                                                          #
#                 PARSING HELPER FUNCTIONS                 #
#          JUMP TO LINE 240 FOR THE ACTUAL SCRIPT          #
#                                                          #
############################################################

# Generated with Argbash
### TEMPLATE
# ARG_OPTIONAL_BOOLEAN([cleanup],[c],[Remove all previous data from elasticsearch and suricata.])
# ARG_OPTIONAL_BOOLEAN([autofp],[a],[Run in autofp mode instead of single mode.])
# ARG_OPTIONAL_SINGLE([set-rulefile],[s],[Set a file with signatures, which will be loaded together with the rules set in the yaml.])
# ARG_OPTIONAL_SINGLE([set-rulefile-exclusive],[S],[Set a file with signatures, which will be loaded exclusively, regardless of the rules set in the yaml.])
# ARG_POSITIONAL_SINGLE([path],[Path to the pcap file to read. If <path> specifies a directory, all files in that directory\nwill be processed in order of modified time maintaining flow state between files.],[])
# ARG_HELP([Pcap reading script through Suricata],[Usage: readpcap.sh [OPTIONS] <path>])
# ARG_POSITIONAL_DOUBLEDASH([])
# ARGBASH_SET_INDENT([  ])
# ARGBASH_GO()
# needed because of Argbash --> m4_ignore([
### START OF CODE GENERATED BY Argbash v2.9.0 one line above ###
# Argbash is a bash code generator used to get arguments parsing right.
# Argbash is FREE SOFTWARE, see https://argbash.io for more info
# Generated online by https://argbash.io/generate


# # When called, the process ends.
# Args:
#   $1: The exit message (print to stderr)
#   $2: The exit code (default is 1)
# if env var _PRINT_HELP is set to 'yes', the usage is print to stderr (prior to $1)
# Example:
#   test -f "$_arg_infile" || _PRINT_HELP=yes die "Can't continue, have to supply file as an argument, got '$_arg_infile'" 4
die()
{
  local _ret="${2:-1}"
  test "${_PRINT_HELP:-no}" = yes && print_help >&2
  echo "$1" >&2
  exit "${_ret}"
}


# Function that evaluates whether a value passed to it begins by a character
# that is a short option of an argument the script knows about.
# This is required in order to support getopts-like short options grouping.
begins_with_short_option()
{
  local first_option all_short_options='casSh'
  first_option="${1:0:1}"
  test "$all_short_options" = "${all_short_options/$first_option/}" && return 1 || return 0
}

# THE DEFAULTS INITIALIZATION - POSITIONALS
# The positional args array has to be reset before the parsing, because it may already be defined
# - for example if this script is sourced by an argbash-powered script.
_positionals=()
# THE DEFAULTS INITIALIZATION - OPTIONALS
_arg_cleanup="off"
_arg_autofp="off"
_arg_set_rulefile=
_arg_set_rulefile_exclusive=


# Function that prints general usage of the script.
# This is useful if users asks for it, or if there is an argument parsing error (unexpected / spurious arguments)
# and it makes sense to remind the user how the script is supposed to be called.
print_help()
{
  printf '%s\n' "Pcap reading script through Suricata"
  printf 'Usage: %s [-c|--(no-)cleanup] [-a|--(no-)autofp] [-s|--set-rulefile <arg>] [-S|--set-rulefile-exclusive <arg>] [-h|--help] [--] <path>\n' "$0"
  printf '\t%s\n' "<path>: Path to the pcap file to read. If <path> specifies a directory, all files in that directory
		will be processed in order of modified time maintaining flow state between files."
  printf '\t%s\n' "-c, --cleanup, --no-cleanup: Remove all previous data from elasticsearch and suricata. (off by default)"
  printf '\t%s\n' "-a, --autofp, --no-autofp: Run in autofp mode instead of single mode. (off by default)"
  printf '\t%s\n' "-s, --set-rulefile: Set a file with signatures, which will be loaded together with the rules set in the yaml. (no default)"
  printf '\t%s\n' "-S, --set-rulefile-exclusive: Set a file with signatures, which will be loaded exclusively, regardless of the rules set in the yaml. (no default)"
  printf '\t%s\n' "-h, --help: Prints help"
  printf '\n%s\n' "Usage: readpcap.sh [OPTIONS] <path>"
}


# The parsing of the command-line
parse_commandline()
{
  _positionals_count=0
  while test $# -gt 0
  do
    _key="$1"
    # If two dashes (i.e. '--') were passed on the command-line,
    # assign the rest of arguments as positional arguments and bail out.
    if test "$_key" = '--'
    then
      shift
      # Handle the case when the double dash is the last argument.
      test $# -gt 0 || break
      _positionals+=("$@")
      _positionals_count=$((_positionals_count + $#))
      shift $(($# - 1))
      _last_positional="$1"
      break
    fi
    case "$_key" in
      # The cleanup argurment doesn't accept a value,
      # we expect the --cleanup or -c, so we watch for them.
      -c|--no-cleanup|--cleanup)
        _arg_cleanup="on"
        test "${1:0:5}" = "--no-" && _arg_cleanup="off"
        ;;
      # We support getopts-style short arguments clustering,
      # so as -c doesn't accept value, other short options may be appended to it, so we watch for -c*.
      # After stripping the leading -c from the argument, we have to make sure
      # that the first character that follows coresponds to a short option.
      -c*)
        _arg_cleanup="on"
        _next="${_key##-c}"
        if test -n "$_next" -a "$_next" != "$_key"
        then
          { begins_with_short_option "$_next" && shift && set -- "-c" "-${_next}" "$@"; } || die "The short option '$_key' can't be decomposed to ${_key:0:2} and -${_key:2}, because ${_key:0:2} doesn't accept value and '-${_key:2:1}' doesn't correspond to a short option."
        fi
        ;;
      # See the comment of option '--cleanup' to see what's going on here - principle is the same.
      -a|--no-autofp|--autofp)
        _arg_autofp="on"
        test "${1:0:5}" = "--no-" && _arg_autofp="off"
        ;;
      # See the comment of option '-c' to see what's going on here - principle is the same.
      -a*)
        _arg_autofp="on"
        _next="${_key##-a}"
        if test -n "$_next" -a "$_next" != "$_key"
        then
          { begins_with_short_option "$_next" && shift && set -- "-a" "-${_next}" "$@"; } || die "The short option '$_key' can't be decomposed to ${_key:0:2} and -${_key:2}, because ${_key:0:2} doesn't accept value and '-${_key:2:1}' doesn't correspond to a short option."
        fi
        ;;
      # We support whitespace as a delimiter between option argument and its value.
      # Therefore, we expect the --set-rulefile or -s value.
      # so we watch for --set-rulefile and -s.
      # Since we know that we got the long or short option,
      # we just reach out for the next argument to get the value.
      -s|--set-rulefile)
        test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
        _arg_set_rulefile="$2"
        shift
        ;;
      # We support the = as a delimiter between option argument and its value.
      # Therefore, we expect --set-rulefile=value, so we watch for --set-rulefile=*
      # For whatever we get, we strip '--set-rulefile=' using the ${var##--set-rulefile=} notation
      # to get the argument value
      --set-rulefile=*)
        _arg_set_rulefile="${_key##--set-rulefile=}"
        ;;
      # We support getopts-style short arguments grouping,
      # so as -s accepts value, we allow it to be appended to it, so we watch for -s*
      # and we strip the leading -s from the argument string using the ${var##-s} notation.
      -s*)
        _arg_set_rulefile="${_key##-s}"
        ;;
      # See the comment of option '--set-rulefile' to see what's going on here - principle is the same.
      -S|--set-rulefile-exclusive)
        test $# -lt 2 && die "Missing value for the optional argument '$_key'." 1
        _arg_set_rulefile_exclusive="$2"
        shift
        ;;
      # See the comment of option '--set-rulefile=' to see what's going on here - principle is the same.
      --set-rulefile-exclusive=*)
        _arg_set_rulefile_exclusive="${_key##--set-rulefile-exclusive=}"
        ;;
      # See the comment of option '-s' to see what's going on here - principle is the same.
      -S*)
        _arg_set_rulefile_exclusive="${_key##-S}"
        ;;
      # See the comment of option '--cleanup' to see what's going on here - principle is the same.
      -h|--help)
        print_help
        exit 0
        ;;
      # See the comment of option '-c' to see what's going on here - principle is the same.
      -h*)
        print_help
        exit 0
        ;;
      *)
        _last_positional="$1"
        _positionals+=("$_last_positional")
        _positionals_count=$((_positionals_count + 1))
        ;;
    esac
    shift
  done
}


# Check that we receive expected amount positional arguments.
# Return 0 if everything is OK, 1 if we have too little arguments
# and 2 if we have too much arguments
handle_passed_args_count()
{
  local _required_args_string="'path'"
  test "${_positionals_count}" -ge 1 || _PRINT_HELP=yes die "FATAL ERROR: Not enough positional arguments - we require exactly 1 (namely: $_required_args_string), but got only ${_positionals_count}." 1
  test "${_positionals_count}" -le 1 || _PRINT_HELP=yes die "FATAL ERROR: There were spurious positional arguments --- we expect exactly 1 (namely: $_required_args_string), but got ${_positionals_count} (the last one was: '${_last_positional}')." 1
}


# Take arguments that we have received, and save them in variables of given names.
# The 'eval' command is needed as the name of target variable is saved into another variable.
assign_positional_args()
{
  local _positional_name _shift_for=$1
  # We have an array of variables to which we want to save positional args values.
  # This array is able to hold array elements as targets.
  # As variables don't contain spaces, they may be held in space-separated string.
  _positional_names="_arg_path "

  shift "$_shift_for"
  for _positional_name in ${_positional_names}
  do
    test $# -gt 0 || break
    eval "$_positional_name=\${1}" || die "Error during argument parsing, possibly an Argbash bug." 1
    shift
  done
}

# Now call all the functions defined above that are needed to get the job done
parse_commandline "$@"
handle_passed_args_count
assign_positional_args 1 "${_positionals[@]}"



##############################################
#         START OF THE ACTUAL SCRIPT         #
##############################################

BASEDIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )/.." &> /dev/null && pwd )"

function Cleanup(){
    ${BASEDIR}/scripts/cleanup.sh
    if [ $? -ne 0 ];then
      exit $?
    fi
}

## TEST IF DOCKER IS ACCESSIBLE TO THIS USER ##
dockerV=$(docker version --format '{{.Server.Version}}' 2>/dev/null)
if [[ -z "$dockerV" ]]; then
  echo "Please run this script as root or with sudo"
  exit 1
fi

# Execute cleanup script if "--cleanup" is passed
if [[ "${_arg_cleanup}" == "on" ]]; then
  Cleanup
fi

# Get absolute direct path from user input
HOST_PATH="$(cd "$(dirname "${_arg_path}")"; pwd)/$(basename "${_arg_path}")"
FILENAME=$( basename ${HOST_PATH} )

if [[ -d ${HOST_PATH} ]]; then
  LOCAL_PATH="/replay"
elif [[ -f ${HOST_PATH} ]]; then
  LOCAL_PATH="/replay/${FILENAME}"
else
  echo "${HOST_PATH} not found"
  die exiting...
fi

# Use autofp mode if "--autofp" is passed
if [[ "${_arg_autofp}" == "on" ]]; then
  MODE="autofp"
else
  MODE="single"
fi

# Test if rules exists
test=$(docker volume inspect selks_suricata-rules 2>&1 1>/dev/null)
if [[ "${test}" == *"Error"* ]]; then
  echo "${test}"
  exit 1
fi

IMAGE="jasonish/suricata:master-amd64"
OPTIONSTRING=""

if [[ -n "${_arg_set_rulefile}" ]]; then
  echo ${_arg_set_rulefile}
  RULE_HOST_PATH="$(cd "$(dirname "${_arg_set_rulefile}")"; pwd)/$(basename "${_arg_set_rulefile}")"
  RULE_FILENAME=$( basename ${RULE_HOST_PATH} )
  if [[ ! -f ${RULE_HOST_PATH} ]]; then
    echo "file '${RULE_HOST_PATH}' does not exist"
    exit 1
  fi
  OPTIONSTRING="${OPTIONSTRING} -s /rules/${RULE_FILENAME}"
  RULE_MOUNT="-v ${RULE_HOST_PATH}:/rules/${RULE_FILENAME}"
  IMAGE="jasonish/suricata:master-amd64-profiling"

elif [[ -n "${_arg_set_rulefile_exclusive}" ]]; then
  echo ${_arg_set_rulefile_exclusive}
  RULE_HOST_PATH="$(cd "$(dirname "${_arg_set_rulefile_exclusive}")"; pwd)/$(basename "${_arg_set_rulefile_exclusive}")"
  RULE_FILENAME=$( basename ${RULE_HOST_PATH} )
  if [[ ! -f ${RULE_HOST_PATH} ]]; then
    echo "file '${RULE_HOST_PATH}' does not exist"
    exit 1
  fi
  OPTIONSTRING="${OPTIONSTRING} -S /rules/${RULE_FILENAME}"
  RULE_MOUNT="-v ${RULE_HOST_PATH}:/rules/${RULE_FILENAME}"
  IMAGE="jasonish/suricata:master-amd64-profiling"
fi

if [[ -n "${_arg_set_rulefile}" && -n "${_arg_set_rulefile_exclusive}" ]]; then
  echo "Only one of '-s' and '-S' can be used at the same time"
  exit 1
fi



docker run --name suricata-replay --rm -it \
--cap-add=net_admin --cap-add=sys_nice \
-v ${BASEDIR}/containers-data/suricata/logs:/var/log/suricata \
-v selks_suricata-rules:/etc/suricata/rules \
-v ${BASEDIR}/containers-data/suricata/etc:/etc/suricata \
-v ${HOST_PATH}:${LOCAL_PATH} \
${RULE_MOUNT} \
--entrypoint /etc/suricata/new_entrypoint.sh \
${IMAGE} -k none -r ${LOCAL_PATH} --runmode ${MODE} -l /var/log/suricata --set sensor-name=${FILENAME} ${OPTIONSTRING}

docker cp ${HOST_PATH} arkime:/readpcap/
docker exec arkime bash -c "\$ARKIMEDIR/bin/capture -r /readpcap/${FILENAME} -t ${FILENAME} >> \$ARKIMEDIR/logs/capture-readpcap.log 2>&1 && rm -rf /readpcap/* "